#include "sm2.h"
#include "Crypto_Types.h"
#include "Crypto_internal.h"
#include "ctr_drbg.h"

#define RANDOM_NUM_LENGTH 32 
#define HASH_SIZE 32
#define SIGNATURE_LENGTH 32
#define CRYPTO_START_SEC_CODE
#include "Crypto_MemMap.h"
#define PUTU32(p,V) \
	((p)[0] = (uint8_t)((V) >> 24), \
	 (p)[1] = (uint8_t)((V) >> 16), \
	 (p)[2] = (uint8_t)((V) >>  8), \
	 (p)[3] = (uint8_t)(V))

/**
 * Calculate user identification Z value
 *
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param a elliptical curve parameter a
 * @param b elliptical curve parameter b
 * @param Gx the x-coordinate of the base point G
 * @param Gy the y-coordinate of the base point G
 * @param Px the x-coordinate of the public key
 * @param Py the y-coordinate of the public key
 * @param Z output Z-value (hash result)
 *
 */
boolean SM2_compute_Z(
	uint8_t* ID,
    uint16_t ID_len,
    uint8_t* a,
    uint8_t* b,
    uint64_t* Gx,
    uint64_t* Gy,
    uint64_t* Px,
    uint64_t* Py,
    uint8_t* Z
);

boolean SM2_HashMessage(const uint8_t* zValue, const uint8_t* message, size_t msgLen, uint8_t* digest);

// 256 random number generation function
FUNC(Std_ReturnType, CRY_CODE) Crypto_256_CtrDrbg(uint8 buf[32]);


/**
 * Calculate SM2 hash message digest
 *
 * @param zValue concatenate constants and hash them yields zValue
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param digest the result of concatenating messages with z-values and hashing them
 */
boolean SM2_HashMessage(const uint8_t* zValue, const uint8_t* message, size_t msgLen, uint8_t* digest)
{
    size_t zValueLen = 32; // zValue 的长度固定为 32 字节
    size_t totalLen = zValueLen + msgLen;


    // 创建一个临时缓冲区存放 zValue 和消息的拼接
    uint8_t* data = (uint8_t*)malloc(totalLen);
    if (data == NULL || zValue == NULL)
    {
        printf("内存分配失败！");
        return false; // 内存分配失败
    }

    (void)ILib_memcpy(data, zValue, zValueLen);
    (void)ILib_memcpy(data + zValueLen, message, msgLen);

    sm3(data, totalLen, digest);

    // 释放临时缓冲区
    free(data);

    return true;
}

/**
 * Calculate user identification Z value
 *
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param a elliptical curve parameter a
 * @param b elliptical curve parameter b
 * @param Gx the x-coordinate of the base point G
 * @param Gy the y-coordinate of the base point G
 * @param Px the x-coordinate of the public key
 * @param Py the y-coordinate of the public key
 * @param Z output Z-value (hash result)
 *
 */
boolean SM2_compute_Z(
    uint8_t* ID,
    uint16_t ID_len,
    uint8_t* a,
    uint8_t* b,
    uint64_t* Gx,
    uint64_t* Gy,
    uint64_t* Px,
    uint64_t* Py,
    uint8_t* Z
)
{
	if (ID == NULL)
	{
		ID = "1234567812345678";
		ID_len = 16;
	}
    sm2_uint256_t _32b_zValue;
    uint8_t _8b_zValue[32];
    size_t buffer_size = 2 + ID_len + 32 * 6;
    uint8_t* buffer = (uint8_t*)malloc(buffer_size);
    if (buffer == NULL) {
        printf("内存分配失败！\n");
        return -1;  // 分配失败
    }
    uint16_t ENTL = ID_len * 8; // ID 长度（比特表示）
    uint32_t pos = 0;
    // Gx,Gy要转换为8位类型
    uint8_t _8b_Gx[32], _8b_Gy[32], _8b_Px[32], _8b_Py[32];
    convert_uint64_to_uint8(Gx, _8b_Gx);
    convert_uint64_to_uint8(Gy, _8b_Gy);
    convert_uint64_to_uint8(Px, _8b_Px);
    convert_uint64_to_uint8(Py, _8b_Py);


    // 拼接 ENTL（2 字节，大端序）
    buffer[pos++] = (ENTL >> 8) & 0xFF;
    buffer[pos++] = ENTL & 0xFF;

    // 拼接 ID
    //memcpy(buffer + pos, ID, ID_len);
    (void)ILib_memcpy(buffer + pos, ID, ID_len);
    pos += ID_len;

    // 拼接 a, b, Gx, Gy, Px, Py
    //memcpy(buffer + pos, a, 32);
    (void)ILib_memcpy(buffer + pos, a, 32);
    pos += 32;

    (void)ILib_memcpy(buffer + pos, b, 32);
    pos += 32;

    (void)ILib_memcpy(buffer + pos, _8b_Gx, 32);
    pos += 32;

    (void)ILib_memcpy(buffer + pos, _8b_Gy, 32);
    pos += 32;

    (void)ILib_memcpy(buffer + pos, _8b_Px, 32);
    pos += 32;

    (void)ILib_memcpy(buffer + pos, _8b_Py, 32);
    pos += 32;

    sm3(buffer, pos, Z);
    return true; // 假设 sm3_final 总是成功
}

void SM2_compute_C3(sm2_uint256_64t x2, uint8_t* message, size_t mlen, sm2_uint256_64t y2, uint8_t *res) {
    size_t buffer_size = 32 + mlen + 32;
    uint8_t* buffer = (uint8_t*)malloc(buffer_size);
    uint8_t _8b_x[32], _8b_y[32];
    convert_uint64_to_uint8(x2, _8b_x);
    convert_uint64_to_uint8(y2, _8b_y);
    size_t pos = 0;
    (void)ILib_memcpy(buffer, _8b_x, 32);
    pos += 32;
    (void)ILib_memcpy(buffer + pos, message, mlen);
    pos += mlen;
    (void)ILib_memcpy(buffer + pos, _8b_y, 32);
    pos += 32;
    sm3(buffer, pos, res);
}

/**
 * @brief
 * @param in 输入数据
 * @param inlen 输入数据长度
 * @param outlen 期望的输出秘钥长度
 * @param out 输出的秘钥数据的存储空间
 * @return
*/
int SM2_kdf(sm2_point_64t point, size_t inlen, size_t outlen, uint8_t* out) {
    uint8_t counter_be[4];
    uint8_t dgst[32];
    uint32_t counter = 1;
    size_t len;
    uint8_t* in;
    int totalLen = 32 * 2+4;
    uint8_t* buffer_Z = (uint8_t*)malloc(totalLen);
    uint8_t _8b_x[32], _8b_y[32];
    convert_uint64_to_uint8(point.x, _8b_x);
    convert_uint64_to_uint8(point.y, _8b_y);
    (void)ILib_memcpy(buffer_Z, _8b_x, 32);
    (void)ILib_memcpy(buffer_Z+32,_8b_y, 32);

    while (outlen) {
        PUTU32(counter_be, counter);
        (void)ILib_memcpy(buffer_Z + 64, counter_be, sizeof(counter_be));

        counter++;
        sm3(buffer_Z, totalLen, dgst);
        //sm3_init(&ctx);
        //sm3_update(&ctx, in, inlen);
        //sm3_update(&ctx, counter_be, sizeof(counter_be));
        //sm3_finish(&ctx, dgst);
        len = outlen < 32 ? outlen : 32;
        memcpy(out, dgst, len);
        out += len;
        outlen -= len;
    }

    // memset(&ctx, 0, sizeof(SM3_CTX));
    memset(dgst, 0, sizeof(dgst));
    free(buffer_Z);
    return 1;
}

FUNC(Std_ReturnType, CRY_CODE) Crypto_256_CtrDrbg(uint8 buf[32])
{
    Crypto_CtrDrbgData ctx;
    Crypto_entropy_context entropy;
    Std_ReturnType ret;
    uint8 key[16];
    (void)ILib_memset(&ctx, 0u, sizeof(Crypto_CtrDrbgData));
    (void)ILib_memset(&entropy, 0u, sizeof(Crypto_entropy_context));
    ctx.entropy_len = 16;
	for (int a = 0; a < 16; a++)
	{
		key[a] = rand() % 255;
	}
	Crypto_ctr_drbg_seed(&ctx, Crypto_entropy_func, &entropy, key, 16);

    // 寰楀埌闅忔満鏁�
    ret = Crypto_ctr_drbg_random(&ctx, buf, 32);  

    (void)ILib_memset(&ctx, 0u, sizeof(Crypto_CtrDrbgData));

    return ret;
}


/**
 * Signature function
 *
 * @param key secret key
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param signature signature value for message
 */
FUNC(Std_ReturnType, CRY_CODE) Crypto_SM2Sign(
    const Crypto_SM2Key_Data* _key,
    uint8_t* message,
    uint32_t msgLen,
    uint8_t* ID,
    uint16_t ID_len,
    Crypto_SM2Signature* signature)
{
    //todo 鏁版嵁鏍煎紡杞崲
	Std_ReturnType ret;
    if (!_key || !message || !signature)
    {
        return E_NOT_OK; /* 参数检查 */
    }

    // Crypto_SM2Key 杞崲涓� Crypto_SM2Key_64
    Crypto_SM2Key_64 key;
    convert_uint8_to_uint64(_key->privateKey,key.privateKey);
    convert_uint8_to_uint64(_key->publicKeyX,key.publicKeyX);
    convert_uint8_to_uint64(_key->publicKeyY,key.publicKeyY);

    uint8_t zValue[32];  

    uint8_t digest[32];  
    uint64_t tr[4], ts[4]; 
    uint8_t tk[32];
    sm2_uint256_64t random_k;

    if (!SM2_compute_Z(ID, ID_len, param_A, param_B, param_G_64.x,
        param_G_64.y, key.publicKeyX, key.publicKeyY, zValue))
    {
        return E_NOT_OK;
    }

    // 拼接 Z 和消息计算摘要
    if (!SM2_HashMessage(zValue, message, msgLen, digest))
    {
        return E_NOT_OK;
    }

    sm2_uint256_64t _64b_digest, _64b_r, _64b_x;
    sm2_uint256_64t mg_r;
    // 随机数 k 和签名计算
    while (1)
    {
    	ret = Crypto_256_CtrDrbg(tk);
		if (ret != E_OK)
			return E_NOT_OK;
		convert_uint8_to_uint64(tk,random_k);
		// random if k >0 and k < n
		if (sm2_uint256t_is_zero(random_k))
		{
			continue;
		}else if(!is_in_N(random_k))
		{
			// mod n
			num2mgmy_modn_64(random_k,param_N,random_k);
			mgmy2num_modn_64(random_k,param_N,random_k);
		}

        // 1. 计算椭圆曲线点 (x1, y1) = [k]G
        sm2_point_64t point;
        point = Crypto_ECC_wNAF_scalar_multiply(mg_G_64, random_k, 4);
        //demo_point_multiply_64(&mg_G_64, random_k, Prime_64, &point);
        mgmy_modp2point_64(&point, Prime_64, &point);
        // 2. 计算 r = (e + x1) mod n
        convert_uint8_to_uint64(digest, _64b_digest);

        mod_add_64(_64b_digest, point.x, param_N_64, _64b_r);

        if (ILib_memcmp(_64b_r, N_256_ZERO_64, sizeof(sm2_uint256_64t)) && compare_value_64(_64b_r, random_k, param_N_64))
        {
            break;
        }
    }

    // 计算 s
    //sm2_uint256_t _32b_privateKey;
    sm2_uint256_64t mg_privateKey, mg_s;
    sm2_uint256_64t tmp1, tmp2;
    //convert_to_uint32(key->privateKey, _32b_privateKey);
    num2mgmy_modn_64(key.privateKey, param_N_64, mg_privateKey);
    num2mgmy_modn_64(_64b_r, param_N_64, mg_r);
    // 1. tmp1 = r * d
    montgomery_multiplication_modn_64_32(mg_r, mg_privateKey, param_N_64, tmp1);
    sm2_uint256_64t _64b_k, mg_k;
    //convert_to_uint32(k, _32b_k);
    num2mgmy_modn_64(random_k, param_N_64, mg_k);
    // 2. tmp2 = k - tmp1
    mod_sub_64(mg_k, tmp1, param_N_64, tmp2);
    // 3. tmp1 = (1 + d)^-1
    sm2_uint256_64t mg_privateKey_add_1, mg_one;
    mod_add_64(mg_privateKey, N_ONE_MOD_n_64, param_N_64, mg_privateKey_add_1);
    mod_inverse_modn_64(mg_privateKey_add_1, param_N_64, tmp1);
    // 4. mg_s = tmp1 * tmp1  
    montgomery_multiplication_modn_64_32(tmp1, tmp2, param_N_64, mg_s);
    mgmy2num_modn_64(mg_s, param_N_64, mg_s);

    sm2_uint256t_copy(_64b_r, tr);
    sm2_uint256t_copy(mg_s, ts);
    convert_uint64_to_uint8(tr, signature->r);
    convert_uint64_to_uint8(ts, signature->s);

    return E_OK;
}

/**
 * Verification function
 *
 * @param key secret key
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param signature signature value for message
 */
FUNC(Std_ReturnType, CRY_CODE) Crypto_SM2Verify(
    const Crypto_SM2Key_Data* _key,
    uint8_t* message,
    uint32_t msgLen,
    uint8_t* ID,
    uint16_t ID_len,
    const Crypto_SM2Signature* signature)
{
	Std_ReturnType ret;
    if (!_key || !message || !signature)
    {
        return E_NOT_OK;
    }

    Crypto_SM2Key_64 key;
    convert_uint8_to_uint64(_key->privateKey,key.privateKey);
    convert_uint8_to_uint64(_key->publicKeyX,key.publicKeyX);
    convert_uint8_to_uint64(_key->publicKeyY,key.publicKeyY);

    sm2_uint256_64t sig_r,sig_s;
    convert_uint8_to_uint64(signature->r,sig_r);
    convert_uint8_to_uint64(signature->s,sig_s);

    uint8_t zValue[THIRTY_TRHREE];
    uint8_t digest[THIRTY_TRHREE];
    uint8_t t[THIRTY_TRHREE], u1[THIRTY_TRHREE], u2[THIRTY_TRHREE];
    
    sm2_point_64t point1, point2, point;
    sm2_uint256_64t _64b_param_N, _64b_t;

     //判断r和s是否在[1, n-1]内
    if (!is_in_N(sig_r) || !is_in_N(sig_s))
        return E_NOT_OK;

    // 计算 Z 值
    if (!SM2_compute_Z(ID, ID_len, param_A, param_B, param_G_64.x,
        param_G_64.y, key.publicKeyX, key.publicKeyY, zValue))
    {
        return E_NOT_OK;
    }

    // 拼接 Z 和消息计算摘要
    if (!SM2_HashMessage(zValue, message, msgLen, digest))
    {
        return E_NOT_OK;
    }

    // 验签步骤
    // t = r + s mod n
    mod_add_64(sig_r, sig_s, param_N_64, _64b_t);
    if (sm2_uint256t_is_zero(_64b_t)) {
        return E_NOT_OK;
    }
    sm2_uint256_64t  _64b_digest, _64b_r_check;
    sm2_point_64t mg_key_point;

    //demo_point_multiply_64(&mg_G_64, signature_64->s, Prime_64, &point1);;
    demo_point_multiply_64(&mg_G_64, sig_s, Prime_64, &point1);
    sm2_point_64t key_point;
    sm2_uint256t_copy(key.publicKeyX, key_point.x);
    sm2_uint256t_copy(key.publicKeyY, key_point.y);
    point_modp2mgmy_64(&key_point, Prime_64, &mg_key_point);
    demo_point_multiply_64(&mg_key_point, _64b_t, Prime_64, &point2);
    Crypto_ECC_jacc_point_add_64(&point1, &point2, Prime_64, &point);
    sm2_point_64t normal_point;
    mgmy_modp2point_64(&point, Prime_64, &normal_point);
    convert_uint8_to_uint64(digest, _64b_digest);
    mod_add_64(_64b_digest, normal_point.x, param_N_64, _64b_r_check);

    if (!ILib_memcmp(sig_r, _64b_r_check, sizeof(_64b_r_check)))
        return E_OK;
    else {
        return E_NOT_OK;
    }
}


FUNC(Std_ReturnType, CRY_CODE) Crypto_SM2Encrypt(const Crypto_SM2Key* key, uint8_t* message, size_t klen, Crypto_SM2Cipher_64 *cipher) {
    Std_ReturnType ret;
    sm2_point_64t mg_C1, C1, mg_point2, point2,S;
    uint8_t* t = (uint8_t*)malloc(klen);
    sm2_uint256_64t h;
    while (true) {
        // 1. 产生随机数
        sm2_uint256_64t random_k;
        // = {
        //     0x59276E27D506861A,
        //     0x16680F3AD9C02DCC,
        //     0xEF3CC1FA3CDBE4CE,
        //     0x6D54B80DEAC1BC21
        // };
        uint8_t tk[32];
        ret = Crypto_256_CtrDrbg(tk);
		if (ret != E_OK)
			return E_NOT_OK;
		convert_uint8_to_uint64(tk,random_k);
        if (sm2_uint256t_is_zero(random_k))
		{
			continue;
		}else if(!is_in_N(random_k))
		{
			// mod n
			num2mgmy_modn_64(random_k,param_N,random_k);
			mgmy2num_modn_64(random_k,param_N,random_k);
		}

        // 2. 计算椭圆曲线点 C1 = [k]G=(x1,y1)
        demo_point_multiply_64(&mg_G_64, random_k, Prime_64, &mg_C1);
        mgmy_modp2point_64(&mg_C1, Prime_64, &C1);
        // 3. 计算椭圆曲线点 S = [h]Pb
        sm2_point_64t key_point, mg_key_point;
        sm2_uint256t_copy(key->publicKeyX, key_point.x);
        sm2_uint256t_copy(key->publicKeyY, key_point.y);
        point_modp2mgmy_64(&key_point, Prime_64, &mg_key_point);
        // 4. 计算椭圆曲线点 point2 = [k]Pb = (x2,y2)
        demo_point_multiply_64(&mg_key_point, random_k, Prime_64, &mg_point2);
        mgmy_modp2point_64(&mg_point2, Prime_64, &point2);
        // 5. 计算KDF(x2||y2,klen)
        SM2_kdf(point2, 64, klen, t);
        if (!sm2_uint256t_is_zero(t)) {
            break;
        }
    }
    // 6. 计算C2 = M + t
    uint8_t* C2 = (uint8_t*)malloc(klen);
    for (size_t i = 0; i < klen; i++) {
        C2[i] = message[i] ^ t[i];
    }
    //printf("\n");
    // 7. 计算C3 = Hash(x2, M ,y2)
    uint8_t C3[32];
    SM2_compute_C3(point2.x, message, klen, point2.y, C3);
    // 8. 输出密文 C = C1 || C2 || C3
    
    (void)ILib_memcpy(cipher->C1.x, C1.x, sizeof(C1.x));
    (void)ILib_memcpy(cipher->C1.y, C1.y, sizeof(C1.y));
    cipher->C2 = C2;
    //(void)ILib_memcpy(cipher->C2, C2, sizeof(C2));
    (void)ILib_memcpy(cipher->C3, C3, sizeof(C3));

    free(t);
    return E_OK;
}

FUNC(Std_ReturnType, CRY_CODE)  Crypto_SM2Decrypt(const Crypto_SM2Key* key, Crypto_SM2Cipher_64 cipher, size_t klen, uint8_t* M_temp) {
    uint8_t* t = (uint8_t*)malloc(klen);
    // 1. 从密文中取出C1， 验证C1是否满足椭圆曲线方程
    if (!is_point_on_curve(&cipher.C1)) {
        return E_NOT_OK;
    }
    // 2. 计算椭圆曲线点S = [h]C1
    //sm2_point_64t S, mg_C1， C1;
    //sm2_uint256_64t h;
    //point_modp2mgmy_64(&cipher.C1, Prime_64, &mg_C1);
    //demo_point_multiply_64(&mg_C1, h, Prime_64, &S);
    //mgmy_modp2point_64(&mg_C1, Prime_64, &C1);
    //if (sm2_uint256t_is_zero(S.x) && sm2_uint256t_is_zero(S.y)) {
    //    return false;
    //}
    
    // 3. 计算[dB]C1 = (x2, y2)
    sm2_point_64t S, mg_C1, C1, mg_S;
    point_modp2mgmy_64(&cipher.C1, Prime_64, &mg_C1);
    demo_point_multiply_64(&mg_C1, key->privateKey, Prime_64, &mg_S);
    mgmy_modp2point_64(&mg_S, Prime_64, &S);
    // 4. 计算KDF（x2||y2, klen）
    SM2_kdf(S, 64, klen, t);
    if (sm2_uint256t_is_zero(t)) {
        return E_NOT_OK;
    }
    // 5. 计算M' = C2 + t
    for (size_t i = 0; i < klen; i++) {
        M_temp[i] = cipher.C2[i] ^ t[i];
    }
    // 6. 计算 u = Hash(x2||M'|y2)
    uint8_t u[32];
    SM2_compute_C3(S.x, M_temp, klen, S.y, u);
    if (ILib_memcmp(u, cipher.C3, sizeof(cipher.C3))) {
        return E_NOT_OK;
    }

    free(t);
    return E_OK;
}
/**
* Print signature
*/
void printSignature_64(const Crypto_SM2Signature_64* sig) {
    printf("--------------------------------------Signature--------------------------------------\n");
    // 打印 r 值
    printf("r: ");
    print_uint64_array(sig->r);
    printf("\n");

    // 打印 s 值
    printf("s: ");
    print_uint64_array(sig->s);
    printf("----------------------------------------------------------------------------------------\n");
    printf("\n");
}
void printCipher_64(const Crypto_SM2Cipher_64* cipher, size_t klen) {
    printf("--------------------------------------Cipher Text--------------------------------------\n");
    // 打印 C1 值
    printf("C1: \n");
    print_uint64_array(cipher->C1.x);
    print_uint64_array(cipher->C1.y);
    printf("\n");

    // 打印 C3 值
    printf("C3: \n");
    print_uint8_array(cipher->C3);
    printf("\n");

    // 打印 C2 值
    printf("C2: \n");
    for (int i = 0; i < klen; i++) {
        printf("%02x ", cipher->C2[i]);
    }
    printf("\n");
    printf("----------------------------------------------------------------------------------------\n");
    printf("\n");


}
void printMessage_64(const uint8_t* message, size_t mlen) {
    printf("-----------------------------------Decrypted Mesessage----------------------------------\n");
    // 打印 r 值
    printf("decrypt-message: ");
    for (int i = 0; i < mlen; i++) {
        printf("%02X ", message[i]);
    }
    printf("\n");
    printf("----------------------------------------------------------------------------------------\n");
    printf("\n");
}

#define CRYPTO_STOP_SEC_CODE
#include "Crypto_MemMap.h"
